package tree;

import java.util.Comparator;

import tree.BinaryTree.Node;

public class AVLTree<E> extends BBST<E> {

	public AVLTree() {
		this(null);
	}

	public AVLTree(Comparator<E> comparator) {
		super(comparator);
	}

	@Override
	protected void afterAdd(Node<E> node) {
		while ((node = node.parent) != null) {
			if (isBalanced(node)) {
				updateHeight(node);
			} else {
				rebalance(node);
				break;
			}
		}
	}
	
	@Override
	protected void afterRemove(tree.BinaryTree.Node<E> node) {
		while ((node = node.parent) != null) {
			if (isBalanced(node)) {
				updateHeight(node);
			} else {
				rebalance(node);
			}
		}
	}
	
	private void rebalance(Node<E> grand) {
		Node<E> parent = ((AVLNode<E>)grand).tallerChild();
		Node<E> node = ((AVLNode<E>)parent).tallerChild();
		
		if (parent.isLeftChild()) {
			if (node.isRightChild()) {
				rotateLeft(parent);
			}
			rotateRight(grand);
		} else {
			if (node.isLeftChild()) {
				rotateRight(parent);
			}
			rotateLeft(grand);
		}
	}
	
	@Override
	protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
		super.afterRotate(grand, parent, child);
		updateHeight(grand);
		updateHeight(parent);
	}
	

	private boolean isBalanced(Node<E> node) {
		return Math.abs(((AVLNode<E>) node).balanceFactor()) <= 1;
	}

	private void updateHeight(Node<E> node) {
		((AVLNode<E>) node).updateHeight();
	}

	@Override
	protected Node<E> createNode(E element, tree.BinaryTree.Node<E> parent) {
		return new AVLNode<>(element, parent);
	}

	private static class AVLNode<E> extends Node<E> {
		int height = 1;

		public AVLNode(E element, Node<E> parent) {
			super(element, parent);
		}

		public void updateHeight() {
			int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
			height = Math.max(leftHeight, rightHeight) + 1;
		}

		public int balanceFactor() {
			int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
			return leftHeight - rightHeight;
		}

		public Node<E> tallerChild() {
			int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
			if (leftHeight > rightHeight) {
				return left;
			} else if (leftHeight < rightHeight) {
				return right;
			} else {
				return this.isLeftChild() ? left : right;
			}
		}

		@Override
		public String toString() {
			StringBuilder sb = new StringBuilder(100);
			sb.append("p").append("(").append(parent == null ? null : parent.element).append(")");
			sb.append("_");
			sb.append(element);
			sb.append("_");
			sb.append("h").append("(").append(height).append(")");
			return sb.toString();
		}
	}
}
